#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AST验证器（C层）
基于抽象语法树的深度验证
检查代码结构完整性和符号引用一致性
"""

import ast
import os
import re
import shutil
import subprocess
import sys
import tempfile
from typing import Dict, List, Set, Optional, Any
from dataclasses import dataclass, field
from .syntax_validator import ValidationResult
from ..utils_module.language_utils import language_to_extensions

_SWIFTC_PATH = shutil.which('swiftc')
_CLANG_PATH = shutil.which('clang')
_KTLINT_PATH = shutil.which('ktlint')
_JAVAC_PATH = shutil.which('javac')


def _detect_darwin_sdk() -> Optional[str]:
    """
    检测macOS/iOS SDK路径
    :return: SDK路径或None
    """
    if sys.platform != 'darwin':
        return None
    xcrun_path = shutil.which('xcrun')
    if not xcrun_path:
        return None
    try:
        result = subprocess.run(
            [xcrun_path, '--show-sdk-path'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True,
            text=True
        )
        sdk_path = result.stdout.strip()
        return sdk_path if sdk_path else None
    except subprocess.SubprocessError:
        return None


_DARWIN_SDK_PATH = _detect_darwin_sdk()


@dataclass
class SymbolScope:
    """符号作用域"""
    name: str
    level: int
    symbols: Set[str] = field(default_factory=set)
    parent: Optional['SymbolScope'] = None
    
    def add_symbol(self, symbol: str):
        self.symbols.add(symbol)
    
    def has_symbol(self, symbol: str) -> bool:
        if symbol in self.symbols:
            return True
        if self.parent:
            return self.parent.has_symbol(symbol)
        return False

class ASTValidator:
    """AST深度验证器"""
    
    def __init__(self, logger=None):
        self.logger = logger

        language_validators: Dict[str, LanguageASTValidator] = {
            'swift': SwiftASTValidator(logger),
            'objc': ObjCASTValidator(logger),
            'kotlin': KotlinASTValidator(logger),
            'java': JavaASTValidator(logger),
        }

        self.validators: Dict[str, LanguageASTValidator] = {}
        for language, validator in language_validators.items():
            for ext in language_to_extensions(language):
                self.validators[ext] = validator

        # Python 单独处理
        self.validators['.py'] = PythonASTValidator(logger)
    
    def validate_file(self, file_path: str, content: str, symbol_map: Dict[str, str] = None) -> ValidationResult:
        """
        验证单个文件的AST结构
        :param file_path: 文件路径
        :param content: 文件内容
        :param symbol_map: 符号映射（可选）
        """
        result = ValidationResult(
            is_valid=True,
            errors=[],
            warnings=[],
            file_path=file_path
        )
        
        # 获取对应的验证器
        validator = None
        for ext, val in self.validators.items():
            if file_path.endswith(ext):
                validator = val
                break
        
        if not validator:
            result.add_warning(f"没有对应的AST验证器: {file_path}")
            return result
        
        try:
            # 执行AST验证
            validator.validate(content, result, symbol_map)
        except Exception as e:
            result.add_error(f"AST验证失败: {str(e)}")
        
        return result
    
    def validate_symbol_consistency(self, original: str, obfuscated: str, 
                                   symbol_map: Dict[str, str], file_path: str) -> ValidationResult:
        """
        验证符号替换的一致性
        :param original: 原始代码
        :param obfuscated: 混淆后代码
        :param symbol_map: 符号映射
        :param file_path: 文件路径
        """
        result = ValidationResult(
            is_valid=True,
            errors=[],
            warnings=[],
            file_path=file_path
        )
        
        # 检查所有符号是否正确替换
        for original_symbol, obfuscated_symbol in symbol_map.items():
            # 计算出现次数
            original_count = len(re.findall(r'\b' + re.escape(original_symbol) + r'\b', original))
            obfuscated_count = len(re.findall(r'\b' + re.escape(obfuscated_symbol) + r'\b', obfuscated))
            
            if original_count > 0 and obfuscated_count == 0:
                result.add_warning(f"符号 {original_symbol} -> {obfuscated_symbol} 可能未正确替换")
            
            # 检查是否有遗漏的原始符号
            remaining_original = len(re.findall(r'\b' + re.escape(original_symbol) + r'\b', obfuscated))
            if remaining_original > 0:
                result.add_error(f"发现未替换的原始符号: {original_symbol} (剩余{remaining_original}处)")
        
        return result

    @staticmethod
    def _extract_string_literals(content: str) -> List[str]:
        """
        提取代码中的字符串字面量
        :param content: 代码内容
        :return: 字符串字面量列表
        """
        string_pattern = re.compile(r'"(?:\\.|[^"\\])*"|\'(?:\\.|[^\'\\])*\'')
        literals: List[str] = []
        for match in string_pattern.finditer(content):
            literal = match.group(0)
            if literal.startswith(('"', "'")) and len(literal) >= 2:
                literals.append(literal[1:-1])
        return literals

    @staticmethod
    def _check_string_literal_leaks(content: str, symbol_map: Dict[str, str], result: ValidationResult, context: str):
        """
        检查字符串字面量中是否泄露了原始符号
        :param content: 代码内容
        :param symbol_map: 符号映射
        :param result: 验证结果对象
        :param context: 上下文描述（如语言）
        """
        if not symbol_map:
            return
        literals = ASTValidator._extract_string_literals(content)
        literal_set = set(literals)
        for original, obfuscated in symbol_map.items():
            if original in literal_set and obfuscated not in literal_set:
                result.add_warning(f"{context} 字符串仍包含原符号 {original}")


class LanguageASTValidator:
    """语言特定的AST验证器基类"""
    
    def __init__(self, logger=None):
        self.logger = logger
    
    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        """验证代码的AST结构"""
        raise NotImplementedError


class PythonASTValidator(LanguageASTValidator):
    """Python AST验证器（使用内置ast模块）"""
    
    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        """使用Python内置AST验证代码"""
        try:
            # 尝试解析AST
            tree = ast.parse(content)
            
            # 验证AST结构
            self._validate_ast_structure(tree, result)
            
            # 如果有符号映射，验证符号使用
            if symbol_map:
                self._validate_symbol_usage(tree, symbol_map, result)
                
        except SyntaxError as e:
            result.add_error(f"Python语法错误: {e.msg}", e.lineno, e.offset)
        except Exception as e:
            result.add_error(f"AST解析失败: {str(e)}")
    
    def _validate_ast_structure(self, tree: ast.AST, result: ValidationResult):
        """验证AST结构完整性"""
        class StructureValidator(ast.NodeVisitor):
            def __init__(self, result):
                self.result = result
                self.scope_stack = []
            
            def visit_FunctionDef(self, node):
                # 检查函数定义
                if not node.name:
                    self.result.add_error(f"函数缺少名称", node.lineno)
                self.generic_visit(node)
            
            def visit_ClassDef(self, node):
                # 检查类定义
                if not node.name:
                    self.result.add_error(f"类缺少名称", node.lineno)
                self.generic_visit(node)
        
        validator = StructureValidator(result)
        validator.visit(tree)
    
    def _validate_symbol_usage(self, tree: ast.AST, symbol_map: Dict[str, str], result: ValidationResult):
        """验证符号使用的一致性"""
        class SymbolValidator(ast.NodeVisitor):
            def __init__(self, symbol_map, result):
                self.symbol_map = symbol_map
                self.result = result
                self.defined_symbols = set()
                self.used_symbols = set()
            
            def visit_Name(self, node):
                if isinstance(node.ctx, ast.Store):
                    self.defined_symbols.add(node.id)
                else:
                    self.used_symbols.add(node.id)
                self.generic_visit(node)
        
        validator = SymbolValidator(symbol_map, result)
        validator.visit(tree)
        
        # 检查未定义的符号使用
        undefined = validator.used_symbols - validator.defined_symbols
        for symbol in undefined:
            if symbol not in __builtins__:
                result.add_warning(f"使用了未定义的符号: {symbol}")


class SwiftASTValidator(LanguageASTValidator):
    """Swift AST验证器（集成SwiftSyntax/编译器检查）"""

    def __init__(self, logger=None):
        super().__init__(logger)
        self.swiftc_path = _SWIFTC_PATH
        self.sdk_path = _DARWIN_SDK_PATH

    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        """验证Swift代码结构"""
        self._validate_with_swiftc(content, result)
        self._basic_validate(content, result, symbol_map)

    def _validate_with_swiftc(self, content: str, result: ValidationResult):
        """使用swiftc进行语法/AST 校验"""
        if not self.swiftc_path:
            return

        tmp_file = None
        try:
            tmp_file = tempfile.NamedTemporaryFile('w', suffix='.swift', delete=False)
            tmp_file.write(content)
            tmp_file.flush()
            tmp_file.close()

            cmd = [self.swiftc_path, '-frontend', '-dump-ast', tmp_file.name]
            if self.sdk_path:
                cmd.extend(['-sdk', self.sdk_path])

            process = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            if process.returncode != 0:
                message = process.stderr.strip() or 'swiftc 验证失败'
                result.add_error(f"SwiftSyntax 验证失败: {message}")
            elif process.stderr.strip():
                if self.logger:
                    self.logger.log_warning(f"swiftc 警告: {process.stderr.strip()}")

        except Exception as exc:
            if self.logger:
                self.logger.log_warning(f"执行swiftc失败: {exc}")
        finally:
            if tmp_file:
                try:
                    os.unlink(tmp_file.name)
                except OSError:
                    pass

    def _basic_validate(self, content: str, result: ValidationResult, symbol_map: Optional[Dict[str, str]]):
        lines = content.split('\n')
        
        # 构建简单的AST结构（作为回退措施）
        self.root_scope = SymbolScope("root", 0)
        self.current_scope = self.root_scope
        
        for line_num, line in enumerate(lines, 1):
            stripped = line.strip()
            
            # 解析类/结构体/枚举声明
            if re.match(r'(class|struct|enum|protocol)\s+(\w+)', stripped):
                match = re.match(r'(class|struct|enum|protocol)\s+(\w+)', stripped)
                if match:
                    type_name = match.group(2)
                    self.current_scope.add_symbol(type_name)
                    # 创建新作用域
                    new_scope = SymbolScope(type_name, self.current_scope.level + 1, parent=self.current_scope)
                    self.current_scope = new_scope
            
            # 解析函数声明
            elif re.match(r'func\s+(\w+)', stripped):
                match = re.match(r'func\s+(\w+)', stripped)
                if match:
                    func_name = match.group(1)
                    self.current_scope.add_symbol(func_name)
            
            # 解析变量声明
            elif re.match(r'(var|let)\s+(\w+)', stripped):
                match = re.match(r'(var|let)\s+(\w+)', stripped)
                if match:
                    var_name = match.group(2)
                    self.current_scope.add_symbol(var_name)
            
            # 检查作用域结束
            elif stripped == '}':
                if self.current_scope.parent:
                    self.current_scope = self.current_scope.parent
        
        # 验证符号引用
        if symbol_map:
            self._validate_swift_symbols(content, symbol_map, result)
    
    def _validate_swift_symbols(self, content: str, symbol_map: Dict[str, str], result: ValidationResult):
        """验证Swift符号使用"""
        # 检查所有混淆后的符号是否在正确的上下文中使用
        for original, obfuscated in symbol_map.items():
            # 检查类型声明
            if re.search(r'(class|struct|enum)\s+' + re.escape(obfuscated), content):
                # 类型声明正确
                pass
            # 检查函数调用
            elif re.search(re.escape(obfuscated) + r'\s*\(', content):
                # 函数调用格式正确
                pass
            # 检查属性访问
            elif re.search(r'\.' + re.escape(obfuscated), content):
                # 属性访问格式正确
                pass

        ASTValidator._check_string_literal_leaks(content, symbol_map, result, 'Swift')


class ObjCASTValidator(LanguageASTValidator):
    """Objective-C AST验证器"""

    def __init__(self, logger=None):
        super().__init__(logger)
        self.clang_path = _CLANG_PATH
        self.sdk_path = _DARWIN_SDK_PATH

    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        self._validate_with_clang(content, result)
        self._basic_validate(content, result, symbol_map)

    def _validate_with_clang(self, content: str, result: ValidationResult):
        if not self.clang_path:
            return

        tmp_file = None
        try:
            suffix = '.m'
            tmp_file = tempfile.NamedTemporaryFile('w', suffix=suffix, delete=False)
            tmp_file.write(content)
            tmp_file.flush()
            tmp_file.close()

            cmd = [self.clang_path, '-x', 'objective-c', '-fsyntax-only', '-fobjc-arc', tmp_file.name]
            if self.sdk_path:
                cmd.extend(['-isysroot', self.sdk_path])

            process = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )

            if process.returncode != 0:
                message = process.stderr.strip() or 'clang 语法检查失败'
                result.add_error(f"Clang 验证失败: {message}")
            elif process.stderr.strip():
                if self.logger:
                    self.logger.log_warning(f"clang 警告: {process.stderr.strip()}")

        except Exception as exc:
            if self.logger:
                self.logger.log_warning(f"执行clang失败: {exc}")
        finally:
            if tmp_file:
                try:
                    os.unlink(tmp_file.name)
                except OSError:
                    pass

    def _basic_validate(self, content: str, result: ValidationResult, symbol_map: Optional[Dict[str, str]]):
        lines = content.split('\n')

        # 跟踪接口和实现
        interfaces = set()
        implementations = set()
        current_class = None

        for line_num, line in enumerate(lines, 1):
            stripped = line.strip()

            # 解析@interface
            if stripped.startswith('@interface'):
                match = re.match(r'@interface\s+(\w+)', stripped)
                if match:
                    class_name = match.group(1)
                    interfaces.add(class_name)
                    current_class = class_name

            # 解析@implementation
            elif stripped.startswith('@implementation'):
                match = re.match(r'@implementation\s+(\w+)', stripped)
                if match:
                    class_name = match.group(1)
                    implementations.add(class_name)
                    current_class = class_name

            # 解析@end
            elif stripped == '@end':
                current_class = None

            # 解析方法声明
            elif re.match(r'^[+-]\s*\([^)]+\)\s*(\w+)', stripped):
                match = re.match(r'^[+-]\s*\([^)]+\)\s*(\w+)', stripped)
                if match and current_class:
                    method_name = match.group(1)
                    if not current_class:
                        result.add_warning(f"方法 {method_name} 不在类定义中", line_num)

        # 检查接口和实现的匹配
        for interface in interfaces:
            if interface not in implementations:
                result.add_warning(f"接口 {interface} 缺少实现")

        # 验证符号映射
        if symbol_map:
            self._validate_objc_symbols(content, symbol_map, result)
    
    def _validate_objc_symbols(self, content: str, symbol_map: Dict[str, str], result: ValidationResult):
        """验证Objective-C符号使用"""
        for original, obfuscated in symbol_map.items():
            # 检查选择器
            if re.search(r'@selector\(' + re.escape(original), content):
                # 应该被替换为混淆后的名称
                if not re.search(r'@selector\(' + re.escape(obfuscated), content):
                    result.add_error(f"选择器未正确替换: {original}")
            
            # 检查消息发送
            if re.search(r'\[\s*\w+\s+' + re.escape(original), content):
                if not re.search(r'\[\s*\w+\s+' + re.escape(obfuscated), content):
                    result.add_warning(f"消息发送可能未正确替换: {original}")

        ASTValidator._check_string_literal_leaks(content, symbol_map, result, 'Objective-C')


class KotlinASTValidator(LanguageASTValidator):
    """Kotlin AST验证器"""
    
    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        """验证Kotlin代码结构"""
        lines = content.split('\n')
        
        # 跟踪类和函数
        classes = set()
        functions = set()
        current_class = None
        
        for line_num, line in enumerate(lines, 1):
            stripped = line.strip()
            
            # 解析类声明
            if re.match(r'(class|interface|object|data\s+class)\s+(\w+)', stripped):
                match = re.match(r'(class|interface|object|data\s+class)\s+(\w+)', stripped)
                if match:
                    class_name = match.group(2)
                    classes.add(class_name)
                    current_class = class_name
            
            # 解析函数声明
            elif re.match(r'fun\s+(\w+)', stripped):
                match = re.match(r'fun\s+(\w+)', stripped)
                if match:
                    func_name = match.group(1)
                    functions.add(func_name)
                    
                    # 检查是否是扩展函数
                    if '.' in stripped.split('fun')[0]:
                        result.add_warning(f"扩展函数 {func_name} 需要特别注意", line_num)
            
            # 解析协程
            elif 'suspend' in stripped and 'fun' in stripped:
                # 挂起函数需要特殊处理
                result.add_warning(f"挂起函数需要协程上下文", line_num)
        
        # 验证符号映射
        if symbol_map:
            self._validate_kotlin_symbols(content, symbol_map, result)
    
    def _validate_kotlin_symbols(self, content: str, symbol_map: Dict[str, str], result: ValidationResult):
        """验证Kotlin符号使用"""
        for original, obfuscated in symbol_map.items():
            # 检查数据类
            if f'data class {original}' in content:
                if f'data class {obfuscated}' not in content:
                    result.add_error(f"数据类未正确替换: {original}")
            
            # 检查伴生对象
            if f'companion object {original}' in content:
                if f'companion object {obfuscated}' not in content:
                    result.add_error(f"伴生对象未正确替换: {original}")

        ASTValidator._check_string_literal_leaks(content, symbol_map, result, 'Kotlin')


class JavaASTValidator(LanguageASTValidator):
    """Java AST验证器"""
    
    def validate(self, content: str, result: ValidationResult, symbol_map: Dict[str, str] = None):
        """验证Java代码结构"""
        lines = content.split('\n')
        
        # 跟踪类、接口和方法
        classes = set()
        interfaces = set()
        methods = set()
        current_class = None
        package_name = None
        
        for line_num, line in enumerate(lines, 1):
            stripped = line.strip()
            
            # 解析包声明
            if stripped.startswith('package '):
                match = re.match(r'package\s+([\w.]+);', stripped)
                if match:
                    package_name = match.group(1)
                else:
                    result.add_error(f"包声明格式错误", line_num)
            
            # 解析类/接口声明
            elif re.match(r'(public\s+)?(class|interface|enum)\s+(\w+)', stripped):
                match = re.match(r'(public\s+)?(class|interface|enum)\s+(\w+)', stripped)
                if match:
                    type_kind = match.group(2)
                    type_name = match.group(3)
                    
                    if type_kind == 'class':
                        classes.add(type_name)
                    elif type_kind == 'interface':
                        interfaces.add(type_name)
                    
                    current_class = type_name
            
            # 解析方法声明
            elif current_class and re.search(r'(public|private|protected|static|final)\s+\w+\s+(\w+)\s*\(', stripped):
                match = re.search(r'(public|private|protected|static|final)\s+\w+\s+(\w+)\s*\(', stripped)
                if match:
                    method_name = match.group(2)
                    methods.add(f"{current_class}.{method_name}")
            
            # 解析注解
            elif stripped.startswith('@'):
                # 验证注解格式
                if not re.match(r'@\w+(\([^)]*\))?', stripped):
                    result.add_warning(f"注解格式可能有误", line_num)
        
        # 验证符号映射
        if symbol_map:
            self._validate_java_symbols(content, symbol_map, result)
    
    def _validate_java_symbols(self, content: str, symbol_map: Dict[str, str], result: ValidationResult):
        """验证Java符号使用"""
        for original, obfuscated in symbol_map.items():
            # 检查类引用
            if f'new {original}' in content:
                if f'new {obfuscated}' not in content:
                    result.add_warning(f"类实例化可能未正确替换: {original}")
            
            # 检查静态调用
            if f'{original}.' in content:
                # 检查是否是静态方法调用
                if re.search(f'{re.escape(original)}\\.\\w+\\(', content):
                    if not re.search(f'{re.escape(obfuscated)}\\.\\w+\\(', content):
                        result.add_warning(f"静态调用可能未正确替换: {original}")
            
            # 检查继承关系
            if f'extends {original}' in content:
                if f'extends {obfuscated}' not in content:
                    result.add_error(f"继承关系未正确替换: {original}")
            
            if f'implements {original}' in content:
                if f'implements {obfuscated}' not in content:
                    result.add_error(f"接口实现未正确替换: {original}")

        ASTValidator._check_string_literal_leaks(content, symbol_map, result, 'Java')
